package com.bigdata.wsr.dataCheck.utils;

import lombok.extern.slf4j.Slf4j;

import javax.sql.DataSource;
import java.sql.*;
import java.util.List;
import java.util.Map;

@Slf4j
public class TableUtil {
    public static void getSourceTableCount(List<String> tableList, Map<String, Integer> tableCountMap, DataSource ds) {
        Connection conn = null;
        PreparedStatement ps = null;
        ResultSet rs = null;
        for (String tableName : tableList) {
            String sql = "select count(*) as data_num from %s";
            sql = String.format(sql, tableName);
            int dataNum = 0;
            try {
                conn = ds.getConnection();
                ps = conn.prepareStatement(sql);
                rs = ps.executeQuery();
                while (rs.next()) {
                    dataNum = rs.getInt("data_num");
                }
            } catch (SQLException e) {
                e.printStackTrace();
            } finally {
                close(rs, ps, conn);
            }
            tableCountMap.put(tableName, dataNum);
        }
    }

    public static void getSinkTableCount(List<String> tableList, Map<String, Integer> tableCountMap, DataSource dorisDs) {
        Connection conn = null;
        PreparedStatement ps = null;
        PreparedStatement psDel = null;
        ResultSet rs = null;
        for (String tableName : tableList) {
            String sql = "select count(*) as data_num from %s where is_delete_doris = '0'";
            sql = String.format(sql, tableName);

            String del = "delete from %s where is_delete_doris = '1'";
            del = String.format(del, tableName);

            int dataNum = 0;
            try {
                conn = dorisDs.getConnection();
                psDel = conn.prepareStatement(del);
                psDel.execute();

                ps = conn.prepareStatement(sql);
                rs = ps.executeQuery();
                while (rs.next()) {
                    dataNum = rs.getInt("data_num");
                }
            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                close(rs, ps, psDel, conn);
            }
            tableCountMap.put(tableName, dataNum);
        }
    }


    public static void close(ResultSet rs, Statement smt, Connection connection) {
        try {
            if (rs != null) {
                rs.close();
            }
            if (smt != null) {
                smt.close();
            }
            if (connection != null && !connection.isClosed()) {
                connection.close();
            }
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }

    public static void close(ResultSet rs, Statement smt, Statement smt2, Connection connection) {
        try {
            if (rs != null) {
                rs.close();
            }
            if (smt != null) {
                smt.close();
            }
            if (smt2 != null) {
                smt2.close();
            }
            if (connection != null && !connection.isClosed()) {
                connection.close();
            }
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }
}
